{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div align=\"right\" style=\"text-align: right\"><i>Peter Norvig<br>Aug 2021</i></div>\n",
    "\n",
    "# Star Battle Puzzles\n",
    "\n",
    "If you thought this notebook was going to be about a video game like StarCraft®, I'm sorry to disapoint you. Instead it is about [**Star Battle puzzles**](https://www.puzzle-star-battle.com/), a genre of Sudoku-like puzzles with these properties:\n",
    "\n",
    "- Like Sudoku, you start with an *N*×*N* board and fill in cells.\n",
    "- Like Sudoku, an *N*×*N* board has 3*N* **units**: *N* columns, *N* rows, and *N* boxes.\n",
    "- Like Sudoku, a well-formed puzzle has exactly one solution.\n",
    "- Unlike Sudoku, the boxes have irregular shapes (not squares) and differing numbers of cells.\n",
    "- Unlike Sudoku, there are only 2 possible values for each cell: star or blank (not 9 digits).\n",
    "- Unlike Sudoku,\n",
    "- The constraints are:\n",
    "  - Each unit (column, row, or box) must have exactly *S* stars.\n",
    "  - Two stars cannot be adjacent (not even diagonally).\n",
    "\n",
    "\n",
    "\n",
    "Here is a board (*S*=2, *N*=10) and its solution from a helpful [Star Battle Rules and Info](https://www.gmpuzzles.com/blog/star-battle-rules-and-info/) page:\n",
    "\n",
    "![](https://www.gmpuzzles.com/images/blog/GM-StarBattle-Ex.png)\n",
    "\n",
    "This “24”-themed puzzle was created by Thomas Snyder for the 24-Hours Puzzle Championship in Hungary in 2012.\n",
    "\n",
    "# Representation\n",
    "\n",
    "Here's how I will represent Star Battle puzzles:\n",
    "\n",
    "- A **cell** is an integer in `range(N * N)`.\n",
    "  - The top row is cells 0, 1, 2, ... left-to-right; the second row is *N*, *N* + 1, *N* + 2, ...; etc.\n",
    "- A **unit** is a set of cells.\n",
    "- A **board** is a named tuple with the attributes:\n",
    "  - `N`: the number of cells on each side of the *N*×*N* board.\n",
    "  - `S`: the number of stars per unit in a solution.\n",
    "  - `units`: a list of all the units on the board.\n",
    "  - `units_for`: a dict where `units_for[c]` is a list of the units that cell `c` is in.\n",
    "  - `pic`: a string; a graphical picture of the differently-shaped boxes.\n",
    "- An intermediate **state** in a search for a solution is a named tuple with the attributes:\n",
    "  - `units`: list of units that have not yet been filled with *S* stars each.\n",
    "  - `stars`: set of cells that have been determined (or guessed) to contain a star.\n",
    "  - `unknowns`: set of cells that might contain either a star or a blank.\n",
    "- A **solution** is a set of cells where the stars go.\n",
    "- A **failure** to find a solution is indicated by `None`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "from functools   import lru_cache\n",
    "from typing      import Optional, Iterator, Set, List\n",
    "\n",
    "Cell  = int\n",
    "Board = namedtuple('Board', 'N, S, units, units_for, pic')\n",
    "State = namedtuple('State', 'units, stars, unknowns')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Representing Boards\n",
    "\n",
    "We can describe the differently-shaped boxes with a **picture** consisting of of *N* lines of *N* non-whitespace characters, where each of *N* distinct characters in the picture corresponds to a box, as in this picture of a 5×5 board with 5 boxes:\n",
    "    \n",
    "    , + ' ' '\n",
    "    , + : : '\n",
    "    , + : : .\n",
    "    , . . . .\n",
    "    , . . . ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_board(S, picture) -> Board:\n",
    "    \"\"\"Create a `Board` from a picture of the boxes.\"\"\"\n",
    "    pic  = ''.join(picture.split()) # eliminate whitespace from picture\n",
    "    N = int(len(pic) ** 0.5) # N is √(number of cells)\n",
    "    assert len(pic) == N * N\n",
    "    side  = range(0, N)\n",
    "    cols  = [{N * r + c for r in side} for c in side]\n",
    "    rows  = [{N * r + c for c in side} for r in side]\n",
    "    boxes = [indexes(pic, ch) for ch in set(pic)] \n",
    "    units = cols + rows + boxes\n",
    "    units_for = {c: [u for u in units if c in u]\n",
    "                 for c in range(N * N)}\n",
    "    return Board(N, S, units, units_for, pic)\n",
    "\n",
    "def indexes(sequence, item) -> Set[int]:\n",
    "    \"\"\"All indexes in sequence where item appears.\"\"\"\n",
    "    return {i for i in range(len(sequence)) if sequence[i] == item}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the 5×5 board, and the  [board](https://krazydad.com/play/starbattle/?kind=10x10&volumeNumber=1&bookNumber=24&puzzleNumber=18)   Barry Hayes shared to introduce me to this type of puzzle, and the \"24\" board. Note that in the \"24\" board the boxes \"2\" and \"t\" form interlocking figure 2s and the boxes \"4\" and \"f\" form interlocking figure 4s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "board5x5 = make_board(1, \"\"\"\n",
    ", + ' ' '\n",
    ", + : : '\n",
    ", + : : .\n",
    ", . . . .\n",
    ", . . . .\n",
    "\"\"\")\n",
    "\n",
    "board1 = make_board(2, \"\"\"\n",
    "` . . . . . | ; ; ;\n",
    "` . . . . . | ; ; ;\n",
    "` ` ` . . . | ; ; ;\n",
    "` , ` . . . . ; ; ;\n",
    ", , , . . + + = = =\n",
    ", , : : + + + + + +\n",
    ", : : ' ' ' ' ' ' +\n",
    ", : : - - ' ' ' ' '\n",
    ", : : : - - - ' ' '\n",
    ", , , - - - - ' ' '\n",
    "\"\"\") \n",
    "\n",
    "board24 = make_board(2, \"\"\"\n",
    ". . . ' ' , , , , ,\n",
    ". 2 2 2 ' 4 , 4 , -\n",
    ". . . 2 ' 4 , 4 - -\n",
    ". 2 2 2 ' 4 4 4 - -\n",
    ". 2 t t t ; f 4 f -\n",
    ". 2 2 2 t ; f 4 f -\n",
    ". . t t t ; f f f -\n",
    ": : t ; ; ; ; ; f -\n",
    ": : t t t : - ; f -\n",
    ": : : : : : - - - -\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving Strategy\n",
    "\n",
    "I will solve puzzles using **depth-first search** with **constraint propagation**. \n",
    "\n",
    "By \"depth-first search\" I mean a procedure that starts from a current state, then creates a new state with a guess that a star should go into some cell *c*, and then tries to solve the rest of the puzzle from the new state. If there is no solution, back up to the old state and guess a different cell for the star. \n",
    "\n",
    "By \"constraint propagation\" I mean that whenever a star is placed, check what implications this has for the rest of the board: what blanks and/or stars *must* be placed in what cells. Constraint propagation may be able to prove that the original guess leads to failure, and it may make future guesses easier. \n",
    "\n",
    "Note that search always creates a new state for each guess (leaving the old state unaltered so that we can back up to it) and constraint propagation always mutates the state (because the changes are inevitable consequences, not guesses). \n",
    "\n",
    "# Constraint Propagation\n",
    "\n",
    "The constraint propagation rules are:\n",
    "      \n",
    "When we **put a star** in a cell:\n",
    "  - Mutate the current state by adding the cell to the set of stars and removing it from the unknowns.\n",
    "  - If any of the cell's units has more than *S* stars, then fail (return `None`).\n",
    "  - Put a blank in each adjacent cell. If you can't, fail.\n",
    "\n",
    "When we **put a blank** in a cell:\n",
    "  - Mutate the current state by removing the cell from the unknowns.\n",
    "  - For each of the cell's units:\n",
    "    - If the number of stars plus unknown cells in the unit is less than *S*, then fail.\n",
    "    - If the number equals *S*, then put stars in the unknown cells. If you can't, fail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def put_star(cell, board, state) -> Optional[State]:\n",
    "    \"\"\"Put a star on the board in the given cell. Mutates state.\n",
    "    Return `None` if it is not possible to place star.\"\"\"\n",
    "    if cell in state.stars:\n",
    "        return state # Already put star in cell\n",
    "    state.unknowns.remove(cell)\n",
    "    state.stars.add(cell)\n",
    "    for unit in board.units_for[cell]:\n",
    "        if count_stars(unit, state) > board.S:\n",
    "            return None\n",
    "    for c in neighbors(cell, board.N):\n",
    "        if c in state.stars or not put_blank(c, board, state):\n",
    "            return None\n",
    "    return state\n",
    "                       \n",
    "def put_blank(cell, board, state) -> Optional[State]:\n",
    "    \"\"\"Put a blank on the board in the given cell. Mutates state.\n",
    "    Return `None` if it is not possible to place blank.\"\"\"\n",
    "    if cell not in state.unknowns:\n",
    "        return state # Already put blank in cell\n",
    "    state.unknowns.remove(cell)\n",
    "    for unit in board.units_for[cell]:\n",
    "        s = count_stars(unit, state)\n",
    "        unknowns = unit & state.unknowns\n",
    "        if s + len(unknowns) < board.S:\n",
    "            return None\n",
    "        if s + len(unknowns) == board.S:\n",
    "            if not all(put_star(c, board, state) for c in unknowns):\n",
    "                return None\n",
    "    return state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_stars(unit, state) -> int: return len(unit & state.stars)\n",
    "\n",
    "@lru_cache()\n",
    "def neighbors(cell, N) -> Set[Cell]:\n",
    "    \"\"\"The set of cells that neighbor a given cell on an N×N board.\"\"\"\n",
    "    dxs = {0, +1 if cell % N != N - 1 else 0, -1 if cell % N != 0 else 0}\n",
    "    dys = {0, +N if cell + N < N ** 2 else 0, -N if cell >= N     else 0}\n",
    "    return {cell + dx + dy \n",
    "            for dy in dys for dx in dxs if dx or dy}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Depth-First Search\n",
    "\n",
    "Here are the two more main functions to do search:\n",
    "1. `solve(board)`: a wrapper function that calls `search` and prints the results.\n",
    "2. `search(board, state)`: where the real work of searching for a solution is done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve(board) -> Set[Cell]:\n",
    "    \"\"\"Call `search` with an initial state; print board and return solution stars.\n",
    "    Raise an error if there is no solution.\"\"\"\n",
    "    stars = next(search(board, initial_state(board))).stars\n",
    "    print_board(board, stars)\n",
    "    return stars\n",
    "\n",
    "def search(board, state) -> Iterator[State]:\n",
    "    \"\"\"Recursive depth-first search for solution(s) to a Star Battle puzzle.\"\"\"\n",
    "    while state.units and count_stars(state.units[0], state) == board.S:\n",
    "        state.units.pop(0) # Discard filled unit(s)\n",
    "    if not state.units:    # Succeed\n",
    "        yield state\n",
    "    else:                  # Guess and recurse\n",
    "        for c in state.units[0] & state.unknowns:\n",
    "            guess_state = put_star(c, board, copy_state(state))\n",
    "            if guess_state is not None:\n",
    "                yield from search(board, guess_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The inputs to `search` are the static `board` and the dynamic `state` of computation, but what the output  should be is less obvious; I considered two choices:\n",
    "1. `search` returns the first state that represents a solution, or `None` if there is no solution.\n",
    "2. `search` is a generator that yields all states that represent a solution.\n",
    "\n",
    "I decided on the second choice for `search`, even though `solve` only looks at the first solution. That way `search` could be used to verify that puzzles are well-formed, for example.`search` works as follows:\n",
    "  - While the state's first unit is already filled with *S* stars, discard that unit and move on to the next one.\n",
    "  - If the state has no remaining unfilled units, then succeed: yield the state.\n",
    "  - Otherwise, guess and recurse: for each unknown cell *c* in the first unit, create a new state with a star in cell *c*, and if placing the star does not lead to constraint-propagation failure then recursively search from that state.\n",
    "\n",
    "Below are the remaining minor functions. *Note:* in `initial_state`, the `units` are `sorted` smallest-unit first, because if, say, a board has *S*=2 and the smallest unit has 3 cells, then you have a 2/3 chance of guessing correctly where the 2 stars should go. With larger units you would be more likely to guess wrong and waste time backing up, so better to do small units first and large units later, when constraint propagation has eliminated some unknown cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def copy_state(s: State): return State(s.units[:], set(s.stars), set(s.unknowns))\n",
    "\n",
    "def initial_state(board) -> State: \n",
    "    \"\"\"The initial state to start the search for a solution to `board`.\"\"\"\n",
    "    return State(units=sorted(board.units, key=len), \n",
    "                 stars=set(), unknowns=set(range(board.N ** 2)))\n",
    "\n",
    "def print_board(board, stars) -> None:\n",
    "    \"\"\"Print a representation of the board before and after placing the stars.\n",
    "    The output is not beautiful, but it is readable.\"\"\"\n",
    "    N = board.N\n",
    "    def row(chars, i) -> str: return ' '.join(chars[i * N:(i + 1) * N])\n",
    "    filled = [('*' if c in stars else ch) for c, ch in enumerate(board.pic)]\n",
    "    for i in range(N):\n",
    "        print(row(board.pic, i), ' ' * N, row(filled, i))\n",
    "    print('Valid' if is_solution(board, stars) else 'Invalid', 'solution')\n",
    "        \n",
    "def is_solution(board, stars) -> bool:\n",
    "    \"\"\"Verify that all units have S stars and that stars are non-adjacent.\"\"\"\n",
    "    return (all(len(stars & unit) == board.S\n",
    "                for unit in board.units) and\n",
    "            all(c1 not in neighbors(c2, board.N)\n",
    "                for c1 in stars for c2 in stars))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solutions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "solve(board5x5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "` . . . . . | ; ; ;            ` . . . . . * ; * ;\n",
      "` . . . . . | ; ; ;            ` * . . * . | ; ; ;\n",
      "` ` ` . . . | ; ; ;            ` ` ` . . . * ; * ;\n",
      "` , ` . . . . ; ; ;            * , * . . . . ; ; ;\n",
      ", , , . . + + = = =            , , , . . + + * = *\n",
      ", , : : + + + + + +            , , : * + * + + + +\n",
      ", : : ' ' ' ' ' ' +            , * : ' ' ' ' ' ' *\n",
      ", : : - - ' ' ' ' '            , : : * - * ' ' ' '\n",
      ", : : : - - - ' ' '            * : : : - - - * ' '\n",
      ", , , - - - - ' ' '            , , * - * - - ' ' '\n",
      "Valid solution\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{6, 8, 11, 14, 26, 28, 30, 32, 47, 49, 53, 55, 61, 69, 73, 75, 80, 87, 92, 94}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve(board1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ". . . ' ' , , , , ,            . . . ' ' , * , , *\n",
      ". 2 2 2 ' 4 , 4 , -            . 2 * 2 * 4 , 4 , -\n",
      ". . . 2 ' 4 , 4 - -            * . . 2 ' 4 , * - -\n",
      ". 2 2 2 ' 4 4 4 - -            . 2 * 2 * 4 4 4 - -\n",
      ". 2 t t t ; f 4 f -            . 2 t t t ; f * f *\n",
      ". 2 2 2 t ; f 4 f -            * 2 2 2 t * f 4 f -\n",
      ". . t t t ; f f f -            . . t * t ; f f * -\n",
      ": : t ; ; ; ; ; f -            : * t ; ; * ; ; f -\n",
      ": : t t t : - ; f -            : : t * t : - ; * -\n",
      ": : : : : : - - - -            : * : : : : * - - -\n",
      "Valid solution\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{6, 9, 12, 14, 20, 27, 32, 34, 47, 49, 50, 55, 63, 68, 71, 75, 83, 88, 91, 96}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solve(board24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 195 ms, sys: 2.98 ms, total: 198 ms\n",
      "Wall time: 197 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for board in (board5x5, board1, board24):\n",
    "    assert next(search(board, initial_state(board))).stars"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
